/*
 * Copyright (c) 2005, 2014 vacoor
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 */
package org.vacoor.nothing.wdb.websocket;

import com.sshtools.j2ssh.SshClient;
import com.sshtools.j2ssh.authentication.AuthenticationProtocolState;
import com.sshtools.j2ssh.authentication.PasswordAuthenticationClient;
import com.sshtools.j2ssh.connection.Channel;
import com.sshtools.j2ssh.connection.ChannelEventAdapter;
import com.sshtools.j2ssh.connection.ChannelOutputStream;
import com.sshtools.j2ssh.session.PseudoTerminal;
import com.sshtools.j2ssh.session.SessionChannelClient;
import com.sshtools.j2ssh.transport.HostKeyVerification;
import com.sshtools.j2ssh.transport.TransportProtocolException;
import com.sshtools.j2ssh.transport.publickey.SshPublicKey;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.websocket.*;
import javax.websocket.server.PathParam;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.nio.charset.Charset;

/**
 * SSH Proxy WebSocket Endpoint
 *
 * @author vacoor
 */
@ServerEndpoint("/sshproxy/{host}/{port}/{username}/{password}")
public class SshProxyEndpoint {
    private static final Logger LOG = LoggerFactory.getLogger(SshProxyEndpoint.class);
    private static final Charset DEFAULT_CHARSET = Charset.forName("UTF-8");

    private SshClient sshClient;
    private SessionChannelClient sshSession;

    /**
     * TODO 不安全仅测试使用
     */
    @OnOpen
    public void onOpen(
            @PathParam("host") String host, @PathParam("port") int port,
            @PathParam("username") String username, @PathParam("password") String password,
            final Session session) throws IOException {
        LOG.debug("Session {} open", session.getId());

        sshClient = new SshClient();
        sshClient.connect(host, port, new YesHostKeyVerificationImpl());

        /*
        PublicKeyAuthenticationClient pkac = new PublicKeyAuthenticationClient();
        pkac.setUsername(username);
        pkac.setKey();
        SshPrivateKeyFile.parse(new byte[0]).toPrivateKey()
        */

        PasswordAuthenticationClient pac = new PasswordAuthenticationClient();
        pac.setUsername(username);
        pac.setPassword(password);

        if (AuthenticationProtocolState.COMPLETE != sshClient.authenticate(pac)) {
            LOG.debug("authentication failed");

            session.close(new CloseReason(CloseReason.CloseCodes.CANNOT_ACCEPT, "authentication failed"));
            return;
        }

        LOG.debug("authentication successful..");

        sshSession = sshClient.openSessionChannel(new ChannelEventDispatcher(session));
        // sshSession.setEnvironmentVariable("TERM", "xterm-color");   // invalid
        // sshSession.setEnvironmentVariable("BASH_SOURCE", "/etc/bashrc:~/.bashprofile:~/.bashrc");
        // mode VT220 解决windows 乱码??
        if (sshSession.requestPseudoTerminal(Pty.of("xterm-color", 80, 24, "")) && sshSession.startShell()) {
            LOG.debug("ssh proxy pty start successful");

            // session.setMaxIdleTimeout(10 * 60 * 1000);
        } else {
            LOG.debug("ssh proxy pty start failed");
            session.close(new CloseReason(CloseReason.CloseCodes.CANNOT_ACCEPT, "ssh proxy pty start failed"));
        }
    }

    @OnMessage
    public void onMessage(String message, Session session) throws IOException {
        LOG.debug("received message: {} from session {}", message, session.getId());

        if (message.startsWith("resize:")) {
            message = message.substring(7);
            String[] size = message.split(",");
            try {
                sshSession.changeTerminalDimensions(Pty.of("xterm-color", Integer.valueOf(size[0]), Integer.valueOf(size[1]), ""));
            } catch (NumberFormatException ignore) {
                // ignore
            }
        } else {
            ChannelOutputStream out = sshSession.getOutputStream();
            out.write(message.getBytes(DEFAULT_CHARSET));
            out.flush();
        }
    }

    @OnClose
    public void onClose(Session session, CloseReason closeReason) {
        LOG.debug("Session {} closed because of {}", session.getId(), closeReason);

        try {
            if (sshSession != null && sshSession.isOpen()) {
                sshSession.close();
                sshSession = null;
            }
        } catch (IOException ignore) {
            // ignore
        }

        if (sshClient != null && sshClient.isConnected()) {
            sshClient.disconnect();
            sshClient = null;
        }
    }

    @OnError
    public void onError(Throwable ex, Session session) {
        LOG.warn("Session exception: {}", ex);

        if (session.isOpen()) {
            try {
                session.close();
            } catch (IOException e) {
                // ignore
            }
        }
    }

    /* ********************
     *
     * ********************/

    private static class ChannelEventDispatcher extends ChannelEventAdapter {
        private final Session session;

        private ChannelEventDispatcher(Session session) {
            this.session = session;
        }

        @Override
        public void onDataReceived(Channel channel, byte[] bytes) {
            try {
                RemoteEndpoint.Basic re = session.getBasicRemote();
                re.sendText(new String(bytes, DEFAULT_CHARSET));
            } catch (IOException ignore) {
                // ignore
            }
        }

        @Override
        public void onChannelClose(Channel channel) {
            try {
                session.close();
            } catch (IOException ignore) {
                // ignore
            }
        }

        @Override
        public void onChannelEOF(Channel channel) {
            try {
                session.close();
            } catch (IOException ignore) {
                // ignore
            }
        }
    }

    private static class YesHostKeyVerificationImpl implements HostKeyVerification {

        @Override
        public boolean verifyHost(String s, SshPublicKey sshPublicKey) throws TransportProtocolException {
            return true;
        }
    }

    private static class Pty implements PseudoTerminal {
        private final String term;
        private final int columns;
        private final int rows;
        private final int width;
        private final int height;
        private final String encodedTerminalModes;

        public static Pty of(String term, int cols, int rows, String encodedTerminalModes) {
            return of(term, cols, rows, 0, 0, encodedTerminalModes);
        }

        public static Pty of(String term, int cols, int rows, int width, int height, String encodedTerminalModes) {
            return new Pty(term, cols, rows, width, height, encodedTerminalModes);
        }

        private Pty(String term, int columns, int rows, int width, int height, String encodedTerminalModes) {
            this.term = term;
            this.columns = columns;
            this.rows = rows;
            this.width = width;
            this.height = height;
            this.encodedTerminalModes = encodedTerminalModes;
        }

        @Override
        public int getColumns() {
            return columns;
        }

        @Override
        public String getEncodedTerminalModes() {
            return encodedTerminalModes;
        }

        @Override
        public int getHeight() {
            return height;
        }

        @Override
        public int getRows() {
            return rows;
        }

        @Override
        public String getTerm() {
            return term;
        }

        @Override
        public int getWidth() {
            return width;
        }
    }
}
